import sys
import os
import cv2
import numpy as np
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import Qt, QThread, pyqtSignal


class ImageProcessor:
    """图像处理核心功能类"""

    def __init__(self):
        self.image = None
        self.processed_image = None
        self.geometry = None  # 形状识别结果
        self.faces = None  # 人脸检测结果
        self.face_database = {}  # 人脸数据库

    # 1. 基础图像处理功能
    def load_image(self, path):
        """加载图像"""
        self.image = cv2.imread(path)
        self.processed_image = self.image.copy()
        return self.convert_to_qt_pixmap(self.processed_image)

    def convert_to_qt_pixmap(self, img):
        """转换为Qt显示格式"""
        if img is not None:
            rgb_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            bytes_per_line = ch * w
            convert_to_Qt_format = QtGui.QImage(rgb_image.data, w, h, bytes_per_line, QtGui.QImage.Format_RGB888)
            p = convert_to_Qt_format.scaled(640, 480, Qt.KeepAspectRatio)
            return QtGui.QPixmap.fromImage(p)
        return None

    def rotate_image(self, angle):
        """旋转图像"""
        if self.processed_image is not None:
            rows, cols = self.processed_image.shape[:2]
            M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
            self.processed_image = cv2.warpAffine(self.processed_image, M, (cols, rows))

    def scale_image(self, scale_factor):
        """缩放图像"""
        if self.processed_image is not None:
            width = int(self.processed_image.shape[1] * scale_factor)
            height = int(self.processed_image.shape[0] * scale_factor)
            self.processed_image = cv2.resize(self.processed_image, (width, height))

    def crop_image(self, x1, y1, x2, y2):
        """裁剪图像"""
        if self.processed_image is not None:
            self.processed_image = self.processed_image[y1:y2, x1:x2]

    def adjust_color(self, brightness, contrast, saturation):
        """调整色彩"""
        if self.processed_image is not None:
            hsv = cv2.cvtColor(self.processed_image, cv2.COLOR_BGR2HSV)
            h, s, v = cv2.split(hsv)

            # 亮度调整（v通道）
            adjusted_v = cv2.addWeighted(v, contrast, v, 0, brightness)
            # 饱和度调整（s通道）
            adjusted_s = np.clip(s * saturation, 0, 255).astype(np.uint8)

            adjusted_hsv = cv2.merge((h, adjusted_s, adjusted_v))
            self.processed_image = cv2.cvtColor(adjusted_hsv, cv2.COLOR_HSV2BGR)

    def smooth_image(self, kernel_size=5):
        """图像平滑处理"""
        if self.processed_image is not None:
            self.processed_image = cv2.GaussianBlur(self.processed_image, (kernel_size, kernel_size), 0)

    def detect_shapes(self):
        """形状识别"""
        if self.processed_image is None:
            return None

        gray = cv2.cvtColor(self.processed_image, cv2.COLOR_BGR2GRAY)
        # 使用自适应阈值提高检测成功率
        thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        shapes = []
        for cnt in contours:
            approx = cv2.approxPolyDP(cnt, 0.02 * cv2.arcLength(cnt, True), True)
            if len(approx) == 3:
                shapes.append(('Triangle', approx))
            elif len(approx) == 4:
                x, y, w, h = cv2.boundingRect(cnt)
                ar = w / float(h)
                shapes.append(('Rectangle' if abs(ar - 1) > 0.1 else 'Square', approx))
            elif len(approx) == 5:
                shapes.append(('Pentagon', approx))
            elif len(approx) == 6:
                shapes.append(('Hexagon', approx))
            else:
                (x, y), radius = cv2.minEnclosingCircle(cnt)
                shapes.append(('Circle', np.int0([[x, y, radius]])))

        self.geometry = shapes
        return self.draw_shape_contours()

    def draw_shape_contours(self):
        """绘制形状轮廓"""
        if self.geometry:
            img = self.processed_image.copy()
            for shape in self.geometry:
                name, approx = shape
                if name == 'Circle':
                    (x, y, r) = approx[0]
                    cv2.circle(img, (int(x), int(y)), int(r), (0, 255, 0), 2)
                    cv2.putText(img, f'Circle', (int(x) - 20, int(y) - 20),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                else:
                    cv2.drawContours(img, [approx], -1, (0, 255, 0), 2)
                    x, y = approx[0][0]
                    cv2.putText(img, name, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            self.processed_image = img
        return self.convert_to_qt_pixmap(self.processed_image)

    # 5. 图像拼接功能
    def stitch_images(self, images):
        """图像拼接"""
        if len(images) < 2:
            return None, "至少需要两张图像进行拼接"

        stitcher = cv2.Stitcher_create() if cv2.__version__.startswith('4') else cv2.createStitcher()
        status, result = stitcher.stitch(images)

        if status == cv2.STITCHER_OK:
            self.processed_image = result
            return self.convert_to_qt_pixmap(result), "拼接成功"
        else:
            error_msgs = {
                cv2.STITCHER_ERR_NEED_MORE_IMGS: "需要更多图像",
                cv2.STITCHER_ERR_HOMOGRAPHY_EST_FAIL: "无法估计单应性矩阵",
                cv2.STITCHER_ERR_CAMERA_PARAMS_ADJUST_FAIL: "相机参数调整失败"
            }
            error_msg = error_msgs.get(status, f"拼接失败，错误代码: {status}")
            return None, error_msg

    # 2. 计算机视觉算法（人脸检测和识别）
    def detect_faces(self):
        """人脸检测（使用OpenCV Haar级联分类器）"""
        if self.processed_image is None:
            return None

        # 确保分类器路径正确
        try:
            face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        except:
            print("警告：无法加载默认人脸分类器，请检查OpenCV安装")
            return None

        gray = cv2.cvtColor(self.processed_image, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.3, 5)

        self.faces = faces
        return self.draw_faces()

    def draw_faces(self):
        """绘制人脸框"""
        img = self.processed_image.copy()

        if self.faces is not None:
            for (x, y, w, h) in self.faces:
                cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)

                # 人脸识别
                if self.face_database:
                    face_roi = self.processed_image[y:y + h, x:x + w]
                    match_name, similarity = self.recognize_face(face_roi)
                    if match_name:
                        cv2.putText(img, f"{match_name} ({similarity:.2f}%)",
                                    (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
                else:
                    cv2.putText(img, 'Face', (x, y - 10),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

        self.processed_image = img
        return self.convert_to_qt_pixmap(img)

    def extract_face_features(self, face_img):
        """提取人脸特征（使用LBPH方法）"""
        if face_img is None or face_img.size == 0:
            return None

        gray = cv2.cvtColor(face_img, cv2.COLOR_BGR2GRAY)
        gray = cv2.resize(gray, (100, 100))  # 统一尺寸
        return gray

    def add_face_to_database(self, name, face_img):
        """将人脸添加到数据库"""
        features = self.extract_face_features(face_img)
        if features is not None:
            self.face_database[name] = features
            return True
        return False

    def recognize_face(self, face_img):
        """人脸识别（与数据库比对）"""
        if not self.face_database or face_img is None or face_img.size == 0:
            return None, 0

        # 提取待识别的人脸特征
        input_features = self.extract_face_features(face_img)
        if input_features is None:
            return None, 0

        # 创建LBPH识别器
        recognizer = cv2.face.LBPHFaceRecognizer_create()

        # 准备训练数据
        faces = []
        labels = []
        label_id = 0
        id_to_name = {}

        for name, features in self.face_database.items():
            faces.append(features)
            labels.append(label_id)
            id_to_name[label_id] = name
            label_id += 1

        # 训练识别器
        if faces and labels:
            recognizer.train(faces, np.array(labels))

            # 预测
            label, confidence = recognizer.predict(input_features)

            # 置信度转换为相似度百分比 (LBPH的confidence值越小越相似)
            similarity = max(0, 100 - confidence)

            if similarity > 40:  # 降低阈值提高识别率
                return id_to_name[label], similarity

        return None, 0

    # 3. 视频处理功能
    def process_video_frame(self, frame):
        """实时视频处理"""
        if frame is not None:
            self.processed_image = frame.copy()
            # 应用形状识别和人脸检测
            self.detect_shapes()
            self.detect_faces()
            return self.convert_to_qt_pixmap(self.processed_image)
        return None


class VideoThread(QThread):
    """视频处理线程"""
    change_pixmap_signal = pyqtSignal(np.ndarray)

    def __init__(self, source):
        super().__init__()
        self.source = source
        self._run_flag = True

    def run(self):
        # 打开视频源
        cap = cv2.VideoCapture(self.source)

        while self._run_flag:
            ret, frame = cap.read()
            if ret:
                self.change_pixmap_signal.emit(frame)
            else:
                break

        # 释放资源
        cap.release()

    def stop(self):
        """停止线程"""
        self._run_flag = False
        self.wait()


class MainWindow(QtWidgets.QMainWindow):
    """主窗口界面类"""

    def __init__(self):
        super().__init__()
        self.processor = ImageProcessor()
        self.initUI()
        self.video_thread = None
        self.image_paths = []  # 存储用于拼接的图像路径

    def initUI(self):
        # 窗口设置
        self.setWindowTitle("PyQt5图像处理平台")
        self.setGeometry(100, 100, 1200, 800)

        # 中央部件
        central_widget = QtWidgets.QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QtWidgets.QVBoxLayout(central_widget)

        # 菜单栏
        self.create_menu()

        # 工具栏
        self.create_toolbar()

        # 分割器布局
        splitter = QtWidgets.QSplitter(Qt.Vertical)

        # 图像/视频显示区域
        display_widget = QtWidgets.QWidget()
        display_layout = QtWidgets.QVBoxLayout(display_widget)

        self.image_label = QtWidgets.QLabel()
        self.image_label.setMinimumSize(640, 480)
        self.image_label.setAlignment(QtCore.Qt.AlignCenter)
        self.image_label.setText("请加载图像或启动视频")
        self.image_label.setStyleSheet("border: 1px solid #ccc;")
        display_layout.addWidget(self.image_label)

        # 状态信息
        self.status_info = QtWidgets.QLabel("就绪")
        self.status_info.setAlignment(QtCore.Qt.AlignLeft)
        self.status_info.setStyleSheet("padding: 5px; background-color: #f0f0f0;")
        display_layout.addWidget(self.status_info)

        splitter.addWidget(display_widget)

        # 参数设置和结果区域
        control_widget = QtWidgets.QWidget()
        control_layout = QtWidgets.QHBoxLayout(control_widget)

        # 左侧控制面板
        left_panel = QtWidgets.QTabWidget()

        # 图像处理选项卡
        image_tab = QtWidgets.QWidget()
        image_layout = QtWidgets.QVBoxLayout(image_tab)

        # 色彩调整
        color_group = QtWidgets.QGroupBox("色彩调整")
        color_layout = QtWidgets.QVBoxLayout(color_group)

        brightness_layout = QtWidgets.QHBoxLayout()
        brightness_layout.addWidget(QtWidgets.QLabel("亮度:"))
        self.brightness_slider = QtWidgets.QSlider(Qt.Horizontal)
        self.brightness_slider.setRange(-100, 100)
        self.brightness_slider.setValue(0)
        self.brightness_slider.valueChanged.connect(self.adjust_color_params)
        brightness_layout.addWidget(self.brightness_slider)
        self.brightness_value = QtWidgets.QLabel("0")
        self.brightness_value.setFixedWidth(30)
        brightness_layout.addWidget(self.brightness_value)
        color_layout.addLayout(brightness_layout)

        contrast_layout = QtWidgets.QHBoxLayout()
        contrast_layout.addWidget(QtWidgets.QLabel("对比度:"))
        self.contrast_slider = QtWidgets.QSlider(Qt.Horizontal)
        self.contrast_slider.setRange(0, 200)
        self.contrast_slider.setValue(100)
        self.contrast_slider.valueChanged.connect(self.adjust_color_params)
        contrast_layout.addWidget(self.contrast_slider)
        self.contrast_value = QtWidgets.QLabel("1.0")
        self.contrast_value.setFixedWidth(30)
        contrast_layout.addWidget(self.contrast_value)
        color_layout.addLayout(contrast_layout)

        saturation_layout = QtWidgets.QHBoxLayout()
        saturation_layout.addWidget(QtWidgets.QLabel("饱和度:"))
        self.saturation_slider = QtWidgets.QSlider(Qt.Horizontal)
        self.saturation_slider.setRange(0, 200)
        self.saturation_slider.setValue(100)
        self.saturation_slider.valueChanged.connect(self.adjust_color_params)
        saturation_layout.addWidget(self.saturation_slider)
        self.saturation_value = QtWidgets.QLabel("1.0")
        self.saturation_value.setFixedWidth(30)
        saturation_layout.addWidget(self.saturation_value)
        color_layout.addLayout(saturation_layout)

        image_layout.addWidget(color_group)

        # 平滑处理
        smooth_group = QtWidgets.QGroupBox("平滑处理")
        smooth_layout = QtWidgets.QVBoxLayout(smooth_group)

        kernel_layout = QtWidgets.QHBoxLayout()
        kernel_layout.addWidget(QtWidgets.QLabel("核大小:"))
        self.kernel_size_spin = QtWidgets.QSpinBox()
        self.kernel_size_spin.setRange(1, 21)
        self.kernel_size_spin.setValue(5)
        self.kernel_size_spin.setSingleStep(2)  # 只能选择奇数
        self.kernel_size_spin.valueChanged.connect(self.apply_smoothing)
        kernel_layout.addWidget(self.kernel_size_spin)
        smooth_layout.addLayout(kernel_layout)

        smooth_btn = QtWidgets.QPushButton("应用平滑")
        smooth_btn.clicked.connect(self.apply_smoothing)
        smooth_layout.addWidget(smooth_btn)

        image_layout.addWidget(smooth_group)

        # 旋转和缩放
        transform_group = QtWidgets.QGroupBox("变换")
        transform_layout = QtWidgets.QVBoxLayout(transform_group)

        rotate_layout = QtWidgets.QHBoxLayout()
        rotate_layout.addWidget(QtWidgets.QLabel("旋转角度:"))
        self.rotate_spin = QtWidgets.QDoubleSpinBox()
        self.rotate_spin.setRange(-360, 360)
        self.rotate_spin.setValue(0)
        self.rotate_spin.setSingleStep(15)
        rotate_layout.addWidget(self.rotate_spin)

        rotate_btn = QtWidgets.QPushButton("旋转")
        rotate_btn.clicked.connect(self.rotate_image)
        rotate_layout.addWidget(rotate_btn)
        transform_layout.addLayout(rotate_layout)

        scale_layout = QtWidgets.QHBoxLayout()
        scale_layout.addWidget(QtWidgets.QLabel("缩放比例:"))
        self.scale_spin = QtWidgets.QDoubleSpinBox()
        self.scale_spin.setRange(0.1, 10)
        self.scale_spin.setValue(1)
        self.scale_spin.setSingleStep(0.1)
        scale_layout.addWidget(self.scale_spin)

        scale_btn = QtWidgets.QPushButton("缩放")
        scale_btn.clicked.connect(self.scale_image)
        scale_layout.addWidget(scale_btn)
        transform_layout.addLayout(scale_layout)

        crop_btn = QtWidgets.QPushButton("裁剪图像")
        crop_btn.clicked.connect(self.crop_image_dialog)
        transform_layout.addWidget(crop_btn)

        image_layout.addWidget(transform_group)

        left_panel.addTab(image_tab, "图像处理")

        # 计算机视觉选项卡
        cv_tab = QtWidgets.QWidget()
        cv_layout = QtWidgets.QVBoxLayout(cv_tab)

        # 形状识别
        shape_group = QtWidgets.QGroupBox("形状识别")
        shape_layout = QtWidgets.QVBoxLayout(shape_group)

        shape_btn = QtWidgets.QPushButton("检测形状")
        shape_btn.clicked.connect(self.detect_shapes)
        shape_layout.addWidget(shape_btn)

        cv_layout.addWidget(shape_group)

        # 人脸识别
        face_group = QtWidgets.QGroupBox("人脸识别")
        face_layout = QtWidgets.QVBoxLayout(face_group)

        face_detect_btn = QtWidgets.QPushButton("检测人脸")
        face_detect_btn.clicked.connect(self.detect_faces)
        face_layout.addWidget(face_detect_btn)

        # 人脸数据库管理
        db_layout = QtWidgets.QHBoxLayout()
        db_layout.addWidget(QtWidgets.QLabel("姓名:"))
        self.face_name_edit = QtWidgets.QLineEdit()
        db_layout.addWidget(self.face_name_edit)

        add_face_btn = QtWidgets.QPushButton("添加到数据库")
        add_face_btn.clicked.connect(self.add_face_to_database)
        db_layout.addWidget(add_face_btn)
        face_layout.addLayout(db_layout)

        # 人脸列表
        self.face_list = QtWidgets.QListWidget()
        face_layout.addWidget(self.face_list)

        clear_db_btn = QtWidgets.QPushButton("清空数据库")
        clear_db_btn.clicked.connect(self.clear_face_database)
        face_layout.addWidget(clear_db_btn)

        cv_layout.addWidget(face_group)

        left_panel.addTab(cv_tab, "计算机视觉")

        # 视频处理选项卡
        video_tab = QtWidgets.QWidget()
        video_layout = QtWidgets.QVBoxLayout(video_tab)

        # 视频控制
        video_group = QtWidgets.QGroupBox("视频控制")
        video_ctrl_layout = QtWidgets.QHBoxLayout(video_group)

        self.play_btn = QtWidgets.QPushButton("播放视频")
        self.play_btn.clicked.connect(self.play_video_dialog)
        video_ctrl_layout.addWidget(self.play_btn)

        self.camera_btn = QtWidgets.QPushButton("打开摄像头")
        self.camera_btn.clicked.connect(self.start_live_camera)
        video_ctrl_layout.addWidget(self.camera_btn)

        self.stop_btn = QtWidgets.QPushButton("停止")
        self.stop_btn.clicked.connect(self.stop_video)
        self.stop_btn.setEnabled(False)
        video_ctrl_layout.addWidget(self.stop_btn)

        video_layout.addWidget(video_group)

        # 视频参数
        video_param_group = QtWidgets.QGroupBox("视频参数")
        video_param_layout = QtWidgets.QVBoxLayout(video_param_group)

        fps_layout = QtWidgets.QHBoxLayout()
        fps_layout.addWidget(QtWidgets.QLabel("帧率:"))
        self.fps_spin = QtWidgets.QSpinBox()
        self.fps_spin.setRange(1, 60)
        self.fps_spin.setValue(30)
        fps_layout.addWidget(self.fps_spin)
        video_param_layout.addLayout(fps_layout)

        video_layout.addWidget(video_param_group)

        left_panel.addTab(video_tab, "视频处理")

        # 图像拼接选项卡
        stitch_tab = QtWidgets.QWidget()
        stitch_layout = QtWidgets.QVBoxLayout(stitch_tab)

        # 图像选择
        stitch_select_group = QtWidgets.QGroupBox("选择图像")
        stitch_select_layout = QtWidgets.QVBoxLayout(stitch_select_group)

        self.stitch_image_list = QtWidgets.QListWidget()
        stitch_select_layout.addWidget(self.stitch_image_list)

        add_stitch_btn = QtWidgets.QPushButton("添加图像")
        add_stitch_btn.clicked.connect(self.add_images_for_stitching)
        stitch_select_layout.addWidget(add_stitch_btn)

        clear_stitch_btn = QtWidgets.QPushButton("清空列表")
        clear_stitch_btn.clicked.connect(self.clear_stitch_images)
        stitch_select_layout.addWidget(clear_stitch_btn)

        stitch_layout.addWidget(stitch_select_group)

        # 拼接按钮
        stitch_btn = QtWidgets.QPushButton("开始拼接")
        stitch_btn.clicked.connect(self.stitch_images)
        stitch_layout.addWidget(stitch_btn)

        left_panel.addTab(stitch_tab, "图像拼接")

        control_layout.addWidget(left_panel)

        # 右侧结果区域
        right_panel = QtWidgets.QTabWidget()

        # 结果信息选项卡
        results_tab = QtWidgets.QWidget()
        results_layout = QtWidgets.QVBoxLayout(results_tab)

        self.results_text = QtWidgets.QTextEdit()
        self.results_text.setReadOnly(True)
        results_layout.addWidget(self.results_text)

        right_panel.addTab(results_tab, "处理结果")

        # 人脸数据库选项卡
        face_db_tab = QtWidgets.QWidget()
        face_db_layout = QtWidgets.QVBoxLayout(face_db_tab)

        self.face_db_grid = QtWidgets.QGridLayout()
        face_db_tab.setLayout(self.face_db_grid)

        right_panel.addTab(face_db_tab, "人脸数据库")

        control_layout.addWidget(right_panel)

        splitter.addWidget(control_widget)

        # 设置分割器比例
        splitter.setSizes([600, 200])

        main_layout.addWidget(splitter)

        # 状态栏
        self.statusBar().showMessage("就绪")

    def create_menu(self):
        menu_bar = self.menuBar()

        # 文件菜单
        file_menu = menu_bar.addMenu("文件")
        load_action = QtWidgets.QAction("加载图像", self)
        load_action.setShortcut("Ctrl+O")
        load_action.triggered.connect(self.load_image_dialog)
        file_menu.addAction(load_action)

        save_action = QtWidgets.QAction("保存图像", self)
        save_action.setShortcut("Ctrl+S")
        save_action.triggered.connect(self.save_image)
        file_menu.addAction(save_action)

        file_menu.addSeparator()

        exit_action = QtWidgets.QAction("退出", self)
        exit_action.setShortcut("Ctrl+Q")
        exit_action.triggered.connect(self.close)
        file_menu.addAction(exit_action)

        # 编辑菜单
        edit_menu = menu_bar.addMenu("编辑")
        rotate_action = QtWidgets.QAction("旋转图像", self)
        rotate_action.setShortcut("Ctrl+R")
        rotate_action.triggered.connect(self.rotate_image)
        edit_menu.addAction(rotate_action)

        scale_action = QtWidgets.QAction("缩放图像", self)
        scale_action.setShortcut("Ctrl+Z")
        scale_action.triggered.connect(self.scale_image)
        edit_menu.addAction(scale_action)

        crop_action = QtWidgets.QAction("裁剪图像", self)
        crop_action.setShortcut("Ctrl+C")
        crop_action.triggered.connect(self.crop_image_dialog)
        edit_menu.addAction(crop_action)

        # 处理菜单
        process_menu = menu_bar.addMenu("处理")
        shape_action = QtWidgets.QAction("形状识别", self)
        shape_action.triggered.connect(self.detect_shapes)
        process_menu.addAction(shape_action)

        face_action = QtWidgets.QAction("人脸检测", self)
        face_action.triggered.connect(self.detect_faces)
        process_menu.addAction(face_action)

        stitch_action = QtWidgets.QAction("图像拼接", self)
        stitch_action.triggered.connect(self.stitch_images)
        process_menu.addAction(stitch_action)

        # 视频菜单
        video_menu = menu_bar.addMenu("视频")
        play_video_action = QtWidgets.QAction("播放视频", self)
        play_video_action.triggered.connect(self.play_video_dialog)
        video_menu.addAction(play_video_action)

        live_camera_action = QtWidgets.QAction("实时摄像头", self)
        live_camera_action.triggered.connect(self.start_live_camera)
        video_menu.addAction(live_camera_action)

        # 帮助菜单
        help_menu = menu_bar.addMenu("帮助")
        about_action = QtWidgets.QAction("关于", self)
        about_action.triggered.connect(self.show_about)
        help_menu.addAction(about_action)

    def create_toolbar(self):
        toolbar = self.addToolBar("工具")
        toolbar.setMovable(False)

        load_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("document-open"), "加载图像", self)
        load_btn.triggered.connect(self.load_image_dialog)
        toolbar.addAction(load_btn)

        save_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("document-save"), "保存图像", self)
        save_btn.triggered.connect(self.save_image)
        toolbar.addAction(save_btn)

        toolbar.addSeparator()

        rotate_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("object-rotate-right"), "旋转", self)
        rotate_btn.triggered.connect(self.rotate_image)
        toolbar.addAction(rotate_btn)

        scale_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("zoom-in"), "缩放", self)
        scale_btn.triggered.connect(self.scale_image)
        toolbar.addAction(scale_btn)

        crop_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("transform-crop"), "裁剪", self)
        crop_btn.triggered.connect(self.crop_image_dialog)
        toolbar.addAction(crop_btn)

        toolbar.addSeparator()

        shape_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("draw-polygon"), "形状识别", self)
        shape_btn.triggered.connect(self.detect_shapes)
        toolbar.addAction(shape_btn)

        face_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("face-smile"), "人脸检测", self)
        face_btn.triggered.connect(self.detect_faces)
        toolbar.addAction(face_btn)

        stitch_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("object-group"), "图像拼接", self)
        stitch_btn.triggered.connect(self.stitch_images)
        toolbar.addAction(stitch_btn)

        toolbar.addSeparator()

        video_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("media-playback-start"), "播放视频", self)
        video_btn.triggered.connect(self.play_video_dialog)
        toolbar.addAction(video_btn)

        camera_btn = QtWidgets.QAction(QtGui.QIcon.fromTheme("camera-photo"), "摄像头", self)
        camera_btn.triggered.connect(self.start_live_camera)
        toolbar.addAction(camera_btn)

    # ---------- 图像操作 ----------
    def load_image_dialog(self):
        path, _ = QtWidgets.QFileDialog.getOpenFileName(
            self, "加载图像", "",
            "图像文件 (*.jpg *.jpeg *.png *.bmp *.gif);;所有文件 (*)"
        )
        if path:
            self.stop_video()  # 停止任何正在运行的视频
            qt_pixmap = self.processor.load_image(path)
            if qt_pixmap:
                self.image_label.setPixmap(qt_pixmap)
                self.status_info.setText(f"已加载图像：{os.path.basename(path)}")
                self.statusBar().showMessage(f"已加载图像：{path}")

    def save_image(self):
        if self.processor.processed_image is not None:
            path, _ = QtWidgets.QFileDialog.getSaveFileName(
                self, "保存图像", "",
                "PNG图像 (*.png);;JPEG图像 (*.jpg);;BMP图像 (*.bmp);;所有文件 (*)"
            )
            if path:
                cv2.imwrite(path, self.processor.processed_image)
                self.status_info.setText(f"图像已保存至：{os.path.basename(path)}")
                self.statusBar().showMessage(f"图像已保存至：{path}")

    def rotate_image(self):
        if self.processor.processed_image is not None:
            angle = self.rotate_spin.value()
            self.processor.rotate_image(angle)
            self.update_display()
            self.status_info.setText(f"图像已旋转 {angle} 度")

    def scale_image(self):
        if self.processor.processed_image is not None:
            factor = self.scale_spin.value()
            self.processor.scale_image(factor)
            self.update_display()
            self.status_info.setText(f"图像已缩放至 {factor * 100}%")

    def crop_image_dialog(self):
        if self.processor.processed_image is not None:
            h, w = self.processor.processed_image.shape[:2]

            dialog = QtWidgets.QDialog(self)
            dialog.setWindowTitle("裁剪图像")
            dialog.setMinimumWidth(300)

            layout = QtWidgets.QFormLayout(dialog)

            x1 = QtWidgets.QSpinBox()
            x1.setRange(0, w - 1)
            y1 = QtWidgets.QSpinBox()
            y1.setRange(0, h - 1)
            x2 = QtWidgets.QSpinBox()
            x2.setRange(1, w)
            x2.setValue(w)
            y2 = QtWidgets.QSpinBox()
            y2.setRange(1, h)
            y2.setValue(h)

            layout.addRow("左边界(X1):", x1)
            layout.addRow("上边界(Y1):", y1)
            layout.addRow("右边界(X2):", x2)
            layout.addRow("下边界(Y2):", y2)

            btn_box = QtWidgets.QDialogButtonBox(
                QtWidgets.QDialogButtonBox.Ok | QtWidgets.QDialogButtonBox.Cancel
            )
            layout.addWidget(btn_box)

            btn_box.accepted.connect(lambda: self.processor.crop_image(x1.value(), y1.value(), x2.value(), y2.value()))
            btn_box.rejected.connect(dialog.reject)

            if dialog.exec_() == QtWidgets.QDialog.Accepted:
                self.update_display()
                self.status_info.setText(f"图像已裁剪")

    def adjust_color_params(self):
        if self.processor.processed_image is not None:
            brightness = self.brightness_slider.value()
            contrast = self.contrast_slider.value() / 100.0  # 转换为0.0-2.0范围
            saturation = self.saturation_slider.value() / 100.0  # 转换为0.0-2.0范围

            # 更新显示值
            self.brightness_value.setText(str(brightness))
            self.contrast_value.setText(f"{contrast:.1f}")
            self.saturation_value.setText(f"{saturation:.1f}")

            self.processor.adjust_color(brightness, contrast, saturation)
            self.update_display()
            self.status_info.setText(f"色彩已调整 (亮度:{brightness}, 对比度:{contrast:.1f}, 饱和度:{saturation:.1f})")

    def apply_smoothing(self):
        if self.processor.processed_image is not None:
            kernel_size = self.kernel_size_spin.value()
            if kernel_size % 2 == 0:  # 确保核大小为奇数
                kernel_size += 1
                self.kernel_size_spin.setValue(kernel_size)

            self.processor.smooth_image(kernel_size)
            self.update_display()
            self.status_info.setText(f"图像已平滑处理 (核大小:{kernel_size})")

    def detect_shapes(self):
        if self.processor.processed_image is not None:
            # 先进行预处理以提高识别成功率
            original = self.processor.processed_image.copy()
            self.processor.smooth_image(3)  # 轻微平滑
            qt_pixmap = self.processor.detect_shapes()
            self.processor.processed_image = original  # 恢复原图，避免修改原图

            if qt_pixmap:
                self.image_label.setPixmap(qt_pixmap)
                if self.processor.geometry:
                    count = len(self.processor.geometry)
                    self.status_info.setText(f"检测到 {count} 个形状")

                    # 更新结果文本
                    results = []
                    shape_counts = {}
                    for shape, _ in self.processor.geometry:
                        shape_counts[shape] = shape_counts.get(shape, 0) + 1

                    results.append("形状识别结果:")
                    for shape, count in shape_counts.items():
                        results.append(f"- {shape}: {count}")

                    self.results_text.setText("\n".join(results))
                else:
                    self.status_info.setText("未检测到形状")
                    QtWidgets.QMessageBox.warning(self, "形状识别", "未检测到明显形状，请尝试调整图像对比度或亮度")

    def detect_faces(self):
        if self.processor.processed_image is not None:
            qt_pixmap = self.processor.detect_faces()
            if qt_pixmap:
                self.image_label.setPixmap(qt_pixmap)
                if self.processor.faces is not None and len(self.processor.faces) > 0:
                    count = len(self.processor.faces)
                    self.status_info.setText(f"检测到 {count} 个人脸")

                    # 更新结果文本
                    results = []
                    results.append("人脸检测结果:")
                    for i, (x, y, w, h) in enumerate(self.processor.faces):
                        results.append(f"- 人脸 #{i + 1}: 位置({x}, {y}), 大小({w}x{h})")

                    self.results_text.setText("\n".join(results))
                else:
                    self.status_info.setText("未检测到人脸")
                    QtWidgets.QMessageBox.warning(self, "人脸检测", "未检测到人脸，请确保图像中有清晰正面人脸")

    # ---------- 图像拼接功能 ----------
    def add_images_for_stitching(self):
        paths, _ = QtWidgets.QFileDialog.getOpenFileNames(
            self, "选择图像", "",
            "图像文件 (*.jpg *.jpeg *.png *.bmp);;所有文件 (*)"
        )

        if paths:
            self.image_paths.extend(paths)
            self.update_stitch_image_list()
            self.status_info.setText(f"已添加 {len(paths)} 张图像用于拼接")

    def update_stitch_image_list(self):
        self.stitch_image_list.clear()
        for path in self.image_paths:
            self.stitch_image_list.addItem(os.path.basename(path))

    def clear_stitch_images(self):
        self.image_paths = []
        self.stitch_image_list.clear()
        self.status_info.setText("图像列表已清空")

    def stitch_images(self):
        if len(self.image_paths) < 2:
            QtWidgets.QMessageBox.warning(self, "警告", "至少需要两张图像进行拼接")
            return

        # 加载所有图像
        images = []
        for path in self.image_paths:
            img = cv2.imread(path)
            if img is not None:
                images.append(img)

        if not images:
            QtWidgets.QMessageBox.warning(self, "错误", "无法加载任何图像")
            return

        # 显示处理中对话框
        progress_dialog = QtWidgets.QProgressDialog("正在拼接图像...", "取消", 0, 0, self)
        progress_dialog.setWindowTitle("请稍候")
        progress_dialog.setWindowModality(Qt.WindowModal)
        progress_dialog.show()

        # 执行拼接
        QtCore.QCoreApplication.processEvents()
        result_pixmap, message = self.processor.stitch_images(images)

        progress_dialog.close()

        if result_pixmap:
            self.image_label.setPixmap(result_pixmap)
            self.status_info.setText("图像拼接成功")
            self.results_text.setText("图像拼接成功!\n\n拼接了 {} 张图像".format(len(images)))
        else:
            QtWidgets.QMessageBox.critical(self, "拼接失败", message)
            self.status_info.setText(f"拼接失败: {message}")

    # ---------- 人脸识别功能 ----------
    def add_face_to_database(self):
        if self.processor.processed_image is None or self.processor.faces is None or len(self.processor.faces) == 0:
            QtWidgets.QMessageBox.warning(self, "警告", "请先加载图像并检测人脸")
            return

        name = self.face_name_edit.text().strip()
        if not name:
            QtWidgets.QMessageBox.warning(self, "警告", "请输入人脸名称")
            return

        # 获取最大的人脸区域
        largest_face = max(self.processor.faces, key=lambda f: f[2] * f[3])
        x, y, w, h = largest_face

        # 提取人脸区域
        face_roi = self.processor.processed_image[y:y + h, x:x + w]

        # 添加到数据库
        if self.processor.add_face_to_database(name, face_roi):
            self.face_list.addItem(name)
            self.status_info.setText(f"已将人脸 '{name}' 添加到数据库")

            # 更新人脸数据库显示
            self.update_face_database_display()
        else:
            QtWidgets.QMessageBox.warning(self, "警告", "无法提取人脸特征，请确保人脸清晰可见")

    def update_face_database_display(self):
        # 清除现有显示
        for i in reversed(range(self.face_db_grid.count())):
            widget = self.face_db_grid.itemAt(i).widget()
            if widget:
                widget.setParent(None)

        # 显示人脸数据库
        row = 0
        col = 0
        max_cols = 3

        for name, features in self.processor.face_database.items():
            # 创建人脸图像标签
            face_label = QtWidgets.QLabel()
            qt_pixmap = self.processor.convert_to_qt_pixmap(features)
            if qt_pixmap:
                face_label.setPixmap(qt_pixmap)

            # 创建名称标签
            name_label = QtWidgets.QLabel(name)
            name_label.setAlignment(Qt.AlignCenter)

            # 创建容器
            container = QtWidgets.QVBoxLayout()
            container.addWidget(face_label)
            container.addWidget(name_label)

            # 添加到网格
            self.face_db_grid.addLayout(container, row, col)

            # 更新行列
            col += 1
            if col >= max_cols:
                col = 0
                row += 1

    def clear_face_database(self):
        self.processor.face_database = {}
        self.face_list.clear()
        self.update_face_database_display()
        self.status_info.setText("人脸数据库已清空")

    # ---------- 视频处理 ----------
    def play_video_dialog(self):
        path, _ = QtWidgets.QFileDialog.getOpenFileName(
            self, "播放视频", "",
            "视频文件 (*.mp4 *.avi *.mkv *.mov *.wmv);;所有文件 (*)"
        )
        if path:
            self.start_video(path)

    def start_live_camera(self):
        self.start_video(0)  # 0表示默认摄像头

    def start_video(self, source):
        self.stop_video()  # 停止任何正在运行的视频

        # 创建视频线程
        self.video_thread = VideoThread(source)
        self.video_thread.change_pixmap_signal.connect(self.update_video_frame)
        self.video_thread.start()

        # 更新UI状态
        self.play_btn.setEnabled(False)
        self.camera_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)

        if isinstance(source, str):
            self.status_info.setText(f"正在播放视频: {os.path.basename(source)}")
        else:
            self.status_info.setText("正在捕获摄像头")

    def stop_video(self):
        if self.video_thread and self.video_thread.isRunning():
            self.video_thread.stop()

        # 更新UI状态
        self.play_btn.setEnabled(True)
        self.camera_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)

        self.status_info.setText("视频已停止")

    def update_video_frame(self, frame):
        # 处理视频帧
        qt_pixmap = self.processor.process_video_frame(frame)
        if qt_pixmap:
            self.image_label.setPixmap(qt_pixmap)

    # ---------- 辅助方法 ----------
    def update_display(self):
        """更新显示区域"""
        if self.processor.processed_image is not None:
            qt_pixmap = self.processor.convert_to_qt_pixmap(self.processor.processed_image)
            if qt_pixmap:
                self.image_label.setPixmap(qt_pixmap)

    def show_about(self):
        """显示关于对话框"""
        QtWidgets.QMessageBox.about(
            self, "关于图像处理平台",
            "PyQt5图像处理平台\n\n"
            "版本: 1.0\n"
            "功能: 图像编辑、形状识别、人脸识别、图像拼接、视频处理\n\n"
            "© 2025 图像处理平台开发团队"
        )

    def closeEvent(self, event):
        """关闭窗口时停止所有视频处理"""
        self.stop_video()
        event.accept()


if __name__ == "__main__":
    # 确保中文显示正常
    import matplotlib

    matplotlib.use('Agg')
    matplotlib.rcParams['font.family'] = ['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC']

    app = QtWidgets.QApplication(sys.argv)
    app.setStyle('Fusion')  # 使用Fusion风格，跨平台一致性更好

    # 设置应用程序样式
    palette = QtGui.QPalette()
    palette.setColor(QtGui.QPalette.Window, QtGui.QColor(53, 53, 53))
    palette.setColor(QtGui.QPalette.WindowText, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.Base, QtGui.QColor(25, 25, 25))
    palette.setColor(QtGui.QPalette.AlternateBase, QtGui.QColor(53, 53, 53))
    palette.setColor(QtGui.QPalette.ToolTipBase, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.ToolTipText, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.Text, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.Button, QtGui.QColor(53, 53, 53))
    palette.setColor(QtGui.QPalette.ButtonText, QtCore.Qt.white)
    palette.setColor(QtGui.QPalette.BrightText, QtCore.Qt.red)
    palette.setColor(QtGui.QPalette.Link, QtGui.QColor(42, 130, 218))
    palette.setColor(QtGui.QPalette.Highlight, QtGui.QColor(42, 130, 218))
    palette.setColor(QtGui.QPalette.HighlightedText, QtCore.Qt.black)
    app.setPalette(palette)

    window = MainWindow()
    window.show()
    sys.exit(app.exec_())